package yang.yu.sunny.utils;

import com.google.common.collect.Sets;
import org.apache.commons.dbcp2.BasicDataSource;
import org.dayatang.configuration.Configuration;
import org.dayatang.configuration.ConfigurationFactory;
import org.dayatang.dbunit.DbUnitUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;

import static java.util.stream.Collectors.toMap;

/**
 * Created by yyang on 15/8/5.
 */
public class DbHelper {

    private Map<String, DataSource> dataSources;

    private String dir = "src/test/resources/";

    @Before
    public void setUp() {
        Set<String> dbNames = Sets.newHashSet("product", "sales");
        dataSources = dbNames.stream()
                .collect(toMap(Function.identity(), this::createDataSource));
    }

    @After
    public void tearDown() throws SQLException {
        dataSources.values().forEach(dataSource -> {
            if (dataSource == null) {
                return;
            }
            if (dataSource instanceof BasicDataSource) {
                try {
                    ((BasicDataSource) dataSource).close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
        });
    }

    private DataSource createDataSource(String db) {
        Configuration configuration = new ConfigurationFactory().fromClasspath("/" + db + "/jdbc.properties");
        BasicDataSource dataSource = new BasicDataSource();
        dataSource.setDriverClassName(configuration.getString("jdbc.driverClassName"));
        dataSource.setUrl(configuration.getString("jdbc.url"));
        dataSource.setUsername(configuration.getString("jdbc.username"));
        dataSource.setPassword(configuration.getString("jdbc.password"));
        return dataSource;
    }

    @Test
    public void exportDTD() {
        dataSources.forEach((name, dataSource) -> {
            DbUnitUtils dbUnit = new DbUnitUtils(dataSource);
            dbUnit.exportDtd(dir + "/" + name, "/dataset.dtd");
        });
    }

    @Test
    public void exportData() {
        dataSources.forEach((name, dataSource) -> {
            DbUnitUtils dbUnit = new DbUnitUtils(dataSource);
            dbUnit.exportDtd(dir + "/" + name, "/sample-data.xml");
        });
    }

    @Test
    public void importData() {
        dataSources.forEach((name, dataSource) -> {
            DbUnitUtils dbUnit = new DbUnitUtils(dataSource);
            dbUnit.importDataFromClasspath("/" + name + "/sample-data.xml");
        });
    }

}
